/*
 * Galaxium Messenger
 * 
 * Copyright (C) 2008 Philippe Durand <draekz@gmail.com>
 * 
 * License: GNU General Public License (GPL)
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * for more details.
 * 
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

using System;
using System.Net;
using System.Net.Sockets;
using System.Text;

namespace Galaxium.Protocol
{
	public class AddressTypeNotSupportedException: ApplicationException
	{
		public byte Number { get; private set; }
		public AddressTypeNotSupportedException (byte number) { Number = number; }
	}
	
	public class CommandNotSupportedException: ApplicationException
	{
		public byte Number { get; private set; }
		public CommandNotSupportedException (byte number) { Number = number; }
	}
	
	public class Socks5Proxy
	{
		private enum Method { Anonymous = 0, UserPass = 2, NoneSupported = 0xFF }
		private enum Command { Connect = 1, Bind = 2, UdpAssociate = 3 }
		private enum AddressType { IPv4 = 1, DomainName = 3, IPv6 = 4 }
		private enum Reply {
			Success, GeneralFailure, NotAllowed, NetworkUnreachable,
			HostUnreachable, ConnectionRefused, TTLExpired,
			CommandNotSupported, AddressTypeNotSupported, Unknown
		}
		
		private static byte[] CreateMethodOffer (params Method[] methods)
		{
			var bytes = new byte [methods.Length + 2];
			bytes [0] = 0x05; // protocol version
			bytes [1] = (byte) methods.Length; // number of methods
			for (int i = 0; i < methods.Length; i++)
				bytes [i + 2] = (byte) methods [i];
			return bytes;
		}
		
		private static Method[] DecodeMethodOffer (byte[] bytes)
		{
			var methods = new Method [bytes [1]];
			for (int i = 0; i < bytes [1]; i++)
				methods [i] = (Method) bytes [i + 2];
			return methods;
		}
		
		
		private static byte[] CreateMethodSelect (Method method)
		{
			return new byte[] { 0x05, (byte) method };
		}
		
		private static Method DecodeMethodSelect (byte[] bytes)
		{
			return (Method) bytes [1];
		}

		
		private static byte[] CreateMessage (byte msg, object addr, ushort port)
		{
			if (!(addr is IPAddress || addr is string))
				throw new ArgumentException ("bind_addr is neither IPAddress, nor String");
			
			AddressType addr_type;
			byte[] addr_bytes;
			
			try {
				var ip = addr is IPAddress ? addr as IPAddress : IPAddress.Parse (addr as string);
				addr_type = ip.AddressFamily == AddressFamily.InterNetwork
					? AddressType.IPv4 : AddressType.IPv6;
				addr_bytes = ip.GetAddressBytes ();
			}
			catch {
				addr_type = AddressType.DomainName;
				var bytes = Encoding.ASCII.GetBytes (addr as string);
				addr_bytes = new byte [bytes.Length + 1];
				addr_bytes [0] = (byte) bytes.Length;
				bytes.CopyTo (addr_bytes, 1);
			}
			
			var result = new byte [addr_bytes.Length + 6];
			result [0] = 0x05;
			result [1] = msg;
			result [2] = 0x00;
			result [3] = (byte) addr_type;
			addr_bytes.CopyTo (result, 4);
			result [result.Length - 2] = (byte) (port >> 8);
			result [result.Length - 1] = (byte) (port & 0xFF);
			return result;
		}
		
		private static byte[] CreateRequest (Command command, object dest_addr, ushort dest_port)
		{
			return CreateMessage ((byte) command, dest_addr, dest_port);
		}
		
		private static byte[] CreateReply (Reply reply, object bind_addr, ushort bind_port)
		{
			return CreateMessage ((byte) reply, bind_addr, bind_port);
		}
		
		
		private static void DecodeMessage (byte[] bytes, out byte msg,
		                                   out object addr, out ushort port)
		{
			msg = bytes [1];
			
			AddressType type;
			try { type = (AddressType) bytes [3]; }
			catch { throw new AddressTypeNotSupportedException (bytes [3]); }
			
			int addr_size;
			
			if (type == AddressType.IPv4 || type == AddressType.IPv6) {
				addr_size = type == AddressType.IPv4 ? 4 : 16;
				var addr_bytes = new byte [addr_size];
				Array.Copy (bytes, 4, addr_bytes, 0, addr_size);
				addr = new IPAddress (addr_bytes);
			}
			else {
				addr_size = bytes [4] + 1;
				var addr_bytes = new byte [bytes [4]];
				Array.Copy (bytes, 5, addr_bytes, 0, addr_bytes.Length);
				addr = Encoding.ASCII.GetString (addr_bytes);
			}
			
			var port_offset = addr_size + 4;
			port = (ushort) ((bytes [port_offset] << 8) | bytes [port_offset + 1]);
		}
		
		private static void DecodeRequest (byte[] bytes, out Command command,
		                                   out object dest_addr, out ushort dest_port)
		{
			byte msg;
			DecodeMessage (bytes, out msg, out dest_addr, out dest_port);
			try { command = (Command) msg; }
			catch { throw new CommandNotSupportedException (msg); }
		}
		
		private static void DecodeReply (byte[] bytes, out Reply reply,
		                                 out object bind_addr, out ushort bind_port)
		{
			byte msg;
			DecodeMessage (bytes, out msg, out bind_addr, out bind_port);
			try { reply = (Reply) msg; }
			catch { reply = Reply.Unknown; }
		}
		
		private static byte[] CreateAuthenticationMessage (string username, string password)
		{
			var username_bytes = Encoding.Default.GetBytes (username);
			var password_bytes = Encoding.Default.GetBytes (password);
			
			var bytes = new byte [username_bytes.Length + password_bytes.Length + 3];
			
			bytes [0] = 0x05;
			bytes [1] = (byte) username_bytes.Length;
			username_bytes.CopyTo (bytes, 2);
			bytes [username_bytes.Length + 2] = (byte) password_bytes.Length;
			password_bytes.CopyTo (bytes, username_bytes.Length + 3);
			return bytes;
		}
		
		private static void DecodeAuthenticationMessage (byte[] bytes, out string username,
		                                                 out string password)
		{
			var uname_length = bytes [1];
			var uname = new byte [uname_length];
			Array.Copy (bytes, 2, uname, 0, uname_length);
			var pass_length = bytes [uname_length + 2];
			var pass = new byte [pass_length];
			Array.Copy (bytes, uname_length + 3, pass, 0, pass_length);
			
			username = Encoding.Default.GetString (uname);
			password = Encoding.Default.GetString (pass);
		}
		
		private static byte[] CreateAuthenticationReply (bool authenticated)
		{
			return new byte[] { (byte) 0x05, authenticated ? (byte) 0x00 : (byte) 0x01 };
		}
		
		private static bool DecodeAuthenticationReply (byte[] bytes)
		{
			return bytes [1] == 0;
		}
		
		private static string[] errorMsgs = {
			"Operation completed successfully.",
			"General SOCKS server failure.",
			"Connection not allowed by ruleset.",
			"Network unreachable.",
			"Host unreachable.",
			"Connection refused.",
			"TTL expired.",
			"Command not supported.",
			"Address type not supported.",
			"Unknown error."
		};
		
		public static Socket ConnectAsClient (string proxyAdress, ushort proxyPort, string destAddress, ushort destPort, string userName, string password)
		{
			Socket socket = null;
			try {
				var response = new byte [262];
				
				IPAddress proxyIP;
				
				try { proxyIP = IPAddress.Parse (proxyAdress); }
				catch { proxyIP = Dns.GetHostEntry (proxyAdress).AddressList [0]; }
				
				socket = new Socket (proxyIP.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
			
				try { socket.Connect (proxyIP, proxyPort); }
				catch (SocketException ex) {
					throw new ProxyConnectionException ("Proxy: " + ex.Message);
				}
				
				socket.Send (CreateMethodOffer (Method.Anonymous, Method.UserPass));
				socket.Receive (response);
			
				var method = DecodeMethodSelect (response);
				
				if (method == Method.NoneSupported)
					throw new ProxyConnectionException ("None of the authentication method was accepted by proxy server.");
				
				if (method == Method.UserPass)
				{
					socket.Send (CreateAuthenticationMessage (userName, password));
					socket.Receive (response);
					
					if (!DecodeAuthenticationReply (response))
						throw new ProxyConnectionException ("Bad Username/Password.");
				}
				
				socket.Send (CreateRequest (Command.Connect, destAddress, destPort));
				socket.Receive (response);
				
				Reply reply; object addr; ushort port;
				DecodeReply (response, out reply, out addr, out port);
			
				if (reply != Reply.Success)
					throw new ProxyConnectionException (errorMsgs [(int) reply]);
	
				Console.WriteLine ("PROXY: Received connect address " + addr.ToString () + " with port " + port.ToString ());
								
				Console.WriteLine ("PROXY: Successfully connected through proxy connection!");
			
				return socket;
			}
			catch {
				try { socket.Close (); }
				catch {}
				throw;
			}
		}
	}
}